#from openopt import NLP, GLP
from scipy import optimize
from numpy import cos, arange, ones, asarray, abs, zeros, sqrt, asscalar, array
import pylab
from string import rjust, ljust, expandtabs
import time
import JT
import Calibration

class Test():
    def __init__(self):            
        #===============================================================================
        #    Data Paths
        #===============================================================================
        fixedImageName = 'C:/Users/bryan/bryan-code/2D3D/vert1/fluoro/ushortim080-LAT.mhd'
        #fixedImageName = 'C:/Users/bryan/bryan-code/trunk/Images/CalibratedDRRImage.mhd'
        inputVolFileName = 'C:/Users/bryan/bryan-code/2D3D/vert1/CT/ucharCTv1r1.mha'
        staFile = 'C:/Users/bryan/bryan-code/2D3D/vert1/CT/ucharCTv1r1.sta'
        calFileLat = 'C:/Users/bryan/bryan-code/trunk/TestData/ext_cal1.txt'
        
        roi = [[210,187],[277,370]]
        
        #===============================================================================
        #    Setup DRR
        #===============================================================================
        xraycam = JT.XRayCamera()
        drr = JT.DRR()
        drr.SetXRayCamera(xraycam)
        drr.SetVolumeFilename(inputVolFileName)
        drr.SetBlackOnWhiteImage()
        drr.InteractionOff()
        
        cal = Calibration.ExternalCalibration()
        cal.LoadConsolidatedCalibration(calFileLat)
        cal.LoadStaFile(staFile)
        
        xraycam.SetExternalCalibration(cal)
        
        drr._volume.SetOriginalTransform(cal._VolumeT)
        drr._volume.SetBodyCenter((3.906, 2.17, -7.378))
        
        self.reg=JT.NewRegistration()
        self.reg.SetFixedFileName(fixedImageName)
        self.reg.SetMovingImageGenerator(drr)
        self.reg.SetRegionOfInterest(roi)
        
        self.T = JT.Testing()


    def run_test(self,n_pose=50,method='gc',debug=True,plot=False):    
    #===============================================================================
    #    Setup Optimizer
    #===============================================================================
    # Select the desired starting pose, between [0,200)
        pose = self.T.GetStartPose(n_pose)
        print "Starting Pose: ", pose
        
        # Choose metric and bounds (step size?) for search
        self.reg.SetImageMetric(method)
        step_length = array([1,1,1,.05,.05,.05])
        # Optimizer specific parameters

        x0 = pose
        f=self.reg.GetValue
        #For NCC metric, T0 should be 0.35999
        T0=.8
        Tf=1e-20
        maxiter=10
        maxeval=200
        maxaccept=maxeval
        dwell=35
        feps=1e-4
        boltzmann=1e-20
        learn_rate=1.0  #Not used in Fast schedule
        quench=1.3
        m=1.25
        n=.1
        schedule='fast'
        
        start_time = time.time()
        (xmin, cov_x, infodict, mesg, ier) = optimize.leastsq(
                                             f, x0, args=(debug,plot),
                                             Dfun = None, full_output = 1, 
                                             col_deriv = 0, ftol = feps, 
                                             xtol = 1e-04, gtol = 0.0, 
                                             maxfev = maxeval, epsfcn = 0.0, 
                                             factor = 1, diag = step_length, 
                                             warning = True)
        end_time = time.time()
        
        #===============================================================================
        #    Display/Save Results
        #===============================================================================
        #(params, fopt, direc, numiter, func_calls, warnflag, allvecs) = retval
        print "Finished"
        print infodict
        print mesg
#        print "Performed %i function evaluations, in %f seconds, best cost function: %f" \
#               % (feval,end_time-start_time,Jmin)
#        
        print "Starting Pose: ", pose
        print "Final Pose   : ", xmin
        
        # Save registration data to file for analysis of tuning parameters
#        log_path = 'C:/Users/bryan/bryan-code/trunk/Results/Reg.log'
#        log = open(log_path,'a')
#        output = (str(Jmin),str(xmin[0]),str(xmin[1]),str(xmin[2]),str(xmin[3]),str(xmin[4]),str(xmin[5]),\
#                  str(x0[0]),str(x0[1]),str(x0[2]),str(x0[3]),str(x0[4]),str(x0[5]),\
#                  str(search_range[0]),str(search_range[1]),str(search_range[2]),str(search_range[3]),str(search_range[4]),str(search_range[5]),\
#                  str(T0),str(Tf),str(maxiter),str(maxeval),str(maxaccept),str(dwell),
#                  str(feps),str(retval),schedule,str(boltzmann),str(feval),method,str(quench),str(m),str(n))
#        
#        output = '\t'.join(output)
#        log.write(output)
#        log.write('\n')
#        log.close()

if __name__ == "__main__":
    T = Test()
#    for i in range(11):
#        T.run_test(n_pose=45+i)
    T.run_test(n_pose=50)
    #pylab.show()


